from PIL import Image, ImageOps, ImageFilter
from split_item import split_image
import pytesseract
import re
import mxnet as mx
from cnocr import CnOcr

ocr = CnOcr()
tesseract_cmd = r'E:\Tesseract-OCR'
remove_re = r"\.|\s|-|…||十|二|\+|`"

def get_label(img, threshold=100, white=True):
    table = []
    for i in range(256):
        if i < threshold:
            table.append(1)
        else:
            table.append(0)
    return img.point(lambda i: table[i] if white else (not table[i]), "1").filter(ImageFilter.MinFilter(3))

def clear_name(name):
    name = re.sub(remove_re, "", name)
    name = name.replace("‖", "I").replace("|", "I")
    # name = name.replace("MHK", "MK").replace("十", "").replace("", "").replace(" ", "")
    return name

def get_items_from_image(image_name):
    input_image = Image.open(image_name)
    img_size = input_image.size
    cropped = input_image.crop((int(img_size[0] * 0.249), int(img_size[1] * 0.12), int(img_size[0] * 0.944), int(img_size[1] * 0.82)))# .resize((int(img_size[0] * 0.8), int(img_size[1] * 0.8)))
    cropped = ImageOps.grayscale(cropped)
    cropped.save("./temp/items.jpg")
    get_label(cropped, white=False, threshold=25).save("./temp/items_bin.jpg")
    item_label = split_image("./temp/items_bin.jpg", "item_color")
    # items_size = (int(cropped.size[0]/4), int(cropped.size[1]/2))
    items_dict = []
    item_count = 0
    for label in item_label:
        try:
            item_count += 1
            print("deal items %d" % item_count)
            item_img = cropped.crop(label)
            item_img.save("./temp/item_%d.jpg" % item_count)
            item_size = item_img.size
            name_img = item_img.crop((0, 0, int(item_size[0]), int(item_size[1]*0.235)))
            amount_img = item_img.crop((int(item_size[0]*0.02), int(item_size[1]*0.8), int(item_size[0]*0.2), int(item_size[1]*0.99)))
            name_img.save("./temp/name_%d.jpg" % (item_count))
            amount_img.save("./temp/amount_%d.jpg" % (item_count))
            name = ocr.ocr("./temp/name_%d.jpg" % (item_count))[0]
            name = clear_name("".join(name))
            amount = pytesseract.image_to_string(amount_img, config="--psm 6 --oem 3 digits")
            if not amount.rstrip():
                amount = ocr.ocr("./temp/amount_%d.jpg" % (item_count))[0]
            amount = "".join(amount).replace('l','1').replace('Q','2')
            items_dict.append((name.rstrip(), amount.rstrip()))
        except Exception as e:
            print(e)
    return items_dict

if __name__ == '__main__':
    items = get_items_from_image("./gain/4.jpg")
    print(items)
